package nmf;

import org.ejml.simple.SimpleMatrix;
import org.ejml.ops.NormOps;

//This is the class we are extending, and we will have to implement
//three abstract methods; initialize, cost, and costString. The costString
//function just gives a name to whatever cost function we decide to use.

public abstract class NMFTool {
   static double epsilon = .01;
   
   public abstract void initialize(SimpleMatrix V, SimpleMatrix W, SimpleMatrix H, int[] trainClass);
   
   public double cost(final SimpleMatrix V, final SimpleMatrix W, final SimpleMatrix H) {
      return NormOps.normP2(V.minus(W.mult(H)).getMatrix());//this was zero
   }
   
   public String costString() {
      // TODO Auto-generated method stub
      return "2-Norm";
   }
   
   public int getNumIterations(final SimpleMatrix V) {
      return 30;
   }
   
   private static void updateW(final SimpleMatrix V, final SimpleMatrix W, final SimpleMatrix H) {
      final SimpleMatrix HT = H.transpose();
      final SimpleMatrix numerator = V.mult(HT);
      final SimpleMatrix denominator = W.mult(H).mult(HT);
      
      for (int i = 0; i < W.numRows(); i++) {
         for (int a = 0; a < W.numCols(); a++) {
            //if (numerator.get(i, a)  != 0)
            W.set(i, a, W.get(i, a) * numerator.get(i, a) / (denominator.get(i, a) + epsilon));
         }
      }
   }
   
   private static void updateH(final SimpleMatrix V, final SimpleMatrix W, final SimpleMatrix H)
   {
      final SimpleMatrix WT = W.transpose();
      final SimpleMatrix numerator = WT.mult(V);
      final SimpleMatrix denominator = WT.mult(W).mult(H);
      
      for (int b = 0; b < H.numRows(); b++) {
         for (int j = 0; j < H.numCols(); j++) {
            if (denominator.get(b, j) != 0) {
               //if (numerator.get(b, j)  != 0)
               H.set(b, j, H.get(b, j) * numerator.get(b, j) / (denominator.get(b, j) + epsilon));
            }
         }
      }
   }
   
   public Results factor(final SimpleMatrix V, final int r, final int[] trainClass) {
      
      int numIters = getNumIterations(V);
      
      final SimpleMatrix W = new SimpleMatrix(new double[V.numRows()][r]), H =
            new SimpleMatrix(new double[r][V.numCols()]);
      
      initialize(V, W, H, trainClass);
      
      do {
         updateH(V, W, H);
         updateW(V, W, H);
      } while (numIters-- > 0);
      
      return new Results(W, H, cost(V, W, H));
   }
   
   public static class Results {
      public Results(final SimpleMatrix W, final SimpleMatrix H, final double error) {
         this.W = W;
         this.H = H;
         this.error = error;
      }
      
      public SimpleMatrix W, H;
      public double error;
   }
}
